# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import json
import logging
import numpy
import os
import time

from coin.experimental.xguo.fruit.base.single_product import SingleProductPriceCalculator
from coin.experimental.xguo.fruit.util.fee import get_fee_rate
from coin.experimental.xguo.fruit.util.util import FillStats
from datatool import OhlcCalculator, MovingWindowData


class SmmPricerConfig(object):
  def __init__(self):
    # numbers
    self.window_period = int(300 * 1e9)
    self.book_level = 5

    self.force_fire_threshold = 3 * 1e-4
    self.force_fire_cooldown = 5 * 1e9
    self.force_fire_cooldown_window = 5 * 1e9
    self.price_pull = 2 * 1e-4
    self.price_push = 2 * 1e-4
    self.force_fire_max_count = 3

    self.max_pos = None
    self.min_pos = None
    self.price_range = None

    # product
    self.trade_symbol = None
    self.ref_symbol = None
    self.get_pos = None

  def __str__(self):
    return json.dumps(self.__dict__, indent=2, default=str)


class SmmPricer(object):
  def __init__(self, config):
    self._logger = logging.getLogger("SmmPricer")
    self._price_calculators = {}
    self._timestamp = None
    self._config = config
    self._rate = MovingWindowData(config.window_period / 1.0e9)
    self._trade_buy_window = MovingWindowData(20)
    self._trade_sell_window = MovingWindowData(20)
    self._ohlc = OhlcCalculator(20, maxlen=60)

    fee_bp = get_fee_rate(config.trade_symbol)[0]
    self._fee_rate = fee_bp * 1e-4
    self._fill_stats = FillStats(config.trade_symbol, fee_bp=fee_bp)
    self._last_valid_feed_ts = None

    products = (config.trade_symbol, config.ref_symbol)
    for prod in products:
      if prod is None:
        self._rpc = None
        continue

      price_calculator = SingleProductPriceCalculator(prod)
      self._price_calculators[prod] = price_calculator

      if prod == config.trade_symbol:
        self._tpc = price_calculator
      elif prod == config.ref_symbol:
        self._rpc = price_calculator
      else:
        raise ValueError()
    self._logger.info("SmmPricer created for %s-%s", config.trade_symbol, config.ref_symbol)
    self._logger.info("config: %s", str(self._config))

  @property
  def trade_symbol(self):
    return self._config.trade_symbol

  @property
  def ref_symbol(self):
    return self._config.ref_symbol

  @property
  def fill_stats(self):
    return self._fill_stats

  def is_large_ticksize(self):
    if self._tpc.price_ticksize_ratio - self._fee_rate > 2e-4:
      return True
    else:
      return False

  def get_pos(self):
    return self._config.get_pos(self.trade_symbol.base)

  def get_zero_centered_pos(self):
    reserve_pos = (self._config.max_pos + self._config.min_pos) / 2.
    return self.get_pos() - reserve_pos

  def get_norm_position(self):  # normalized to [-1, 1]
    pos_range = self._config.max_pos - self._config.min_pos

    if abs(pos_range) < 1e-5:
      return 0.0

    n_pos = 2 * (self.get_pos() - self._config.min_pos) / pos_range - 1
    return numpy.clip(n_pos, -1, 1)

  def is_ready(self):
    return len(self._ohlc.data) > 1 and self._fill_stats.is_initialized()

  def update_tick(self):
    self._tpc.renew_tick()
    if self._rpc is not None:
      self._rpc.renew_tick()

  def update(self, product, book):
    self._timestamp = book.timestamp
    self._price_calculators[product].update(book)
    if self._fill_stats.product == product:
      if not self._fill_stats.is_initialized():
        self._fill_stats.set_init_pos(self.get_zero_centered_pos())
      self._fill_stats.update_book(book)

    if self._tpc.is_ready() and product == self.trade_symbol:
      self._ohlc.update(self._timestamp, self._tpc.midp)

    if self._tpc.is_ready() and self._rpc and self._rpc.is_ready():
      implied_rate = self._tpc.midp / self._rpc.midp
      self._rate.update(self._timestamp, implied_rate)

    if self._tpc.product == product:
      self._trade_buy_window.update(self._timestamp, self._tpc.bid0)
      self._trade_sell_window.update(self._timestamp, self._tpc.ask0)

  def update_fill(self, fill_price, fill_qty, order):
    self._fill_stats.update(fill_price, fill_qty, order)

  def calc_edge(self, norm_pos, full_edge):
    full_edge = min(full_edge, 20e-4)  # full edge not too big
    pos = abs(norm_pos)
    # edge must cover fee_rate
    e2 = max(self._fee_rate, full_edge)
    e1 = (1 - pos) * e2 + max(self._fee_rate, full_edge * 0.7) * pos

    assert e1 >= 0, e1
    assert e2 >= 0, e2

    # return buy_edge, sell_edge
    if norm_pos > 0:
      return e2, e1
    else:
      return e1, e2

  def gen_order(self):
    if not self._tpc.is_ready():
      if self._last_valid_feed_ts:
        if (time.time() - self._last_valid_feed_ts/1e9) > 120:
          self._logger.info('invalid book for too long time! restart!')
          # os._exit(0)
      return

    if self.ref_symbol and not self._rpc.is_ready():
      return

    self._last_valid_feed_ts = self._timestamp
    all_orders = {'product': str(self.trade_symbol)}

    if self.is_large_ticksize():
      trade_edge = 0
      trade_buy_price = self._tpc.bid0
      trade_sell_price = self._tpc.ask0
    else:
      trade_edge = numpy.percentile([abs(v.rtn) for v in self._ohlc.data], 90)
      last_trade_edge = self._ohlc.data[-1].rtn
      trade_edge = max(trade_edge, last_trade_edge, self._fee_rate, 1e-4)
      trade_buy_price = numpy.median(list(self._trade_buy_window)) * (1 - trade_edge)
      trade_sell_price = numpy.median(list(self._trade_sell_window)) * (1 + trade_edge)

    norm_pos = self.get_norm_position()
    ref_buy_price = None
    ref_sell_price = None
    e1 = None
    e2 = None
    if self.ref_symbol:
      lower, median, upper = numpy.percentile(list(self._rate), [1, 50, 99])
      ref_buy_price = self._rpc.midp * lower
      ref_sell_price = self._rpc.midp * upper

    if self._config.ref_symbol:
      buy_price = min(trade_buy_price, ref_buy_price)
      sell_price = max(trade_sell_price, ref_sell_price)
    else:
      buy_price = trade_buy_price
      sell_price = trade_sell_price

    fill_buy_price = self._fill_stats.get_lowest_fill_buy_price()
    fill_sell_price = self._fill_stats.get_highest_fill_sell_price()

    fill_elapsed_time = self._fill_stats.get_last_fill_elapsed_time()
    if fill_elapsed_time and fill_elapsed_time > 300 * 1e9:
      if norm_pos > 1e-3:  # if max_pos
        # divided by self._tpc.bid0, changed to ask0
        r = (self._fill_stats.get_highest_fill_buy_price() or 1e-8) / self._tpc.ask0
        if r > 1.001:  # if price has dip a lot.
          self._fill_stats.reset_fill_info()
          return
      elif norm_pos < -1e-3:  # if min_pos
        # was self._tpc.ask0, changed to bid0
        r = self._tpc.bid0 / (self._fill_stats.get_lowest_fill_sell_price() or 1e9)
        if r > 1.001:  # if price has gone up a lot
          self._fill_stats.reset_fill_info()
          return

    if fill_sell_price:
      price = fill_sell_price * (1 - self._fee_rate - 1e-4)
      price = self._tpc.get_prev_price(price)
      buy_price = min(buy_price, price)

    if fill_buy_price:
      price = fill_buy_price * (1 + self._fee_rate + 1e-4)
      price = self._tpc.get_next_price(price)
      sell_price = max(sell_price, price)

    if self.is_large_ticksize():
      pass_buy_price = self._tpc.round_price(buy_price)
      pass_sell_price = self._tpc.round_price(sell_price)
    else:
      pass_buy_price = self._tpc.get_prev_price(buy_price)
      pass_sell_price = self._tpc.get_next_price(sell_price)

    buy_price_pull = pass_buy_price * (1 + self._config.price_pull)
    sell_price_pull = pass_sell_price * (1 - self._config.price_pull)
    buy_price_push = pass_buy_price * (1 - self._config.price_push)
    sell_price_push = pass_sell_price * (1 + self._config.price_push)

    all_orders['mm'] = dict(
        pass_sell_price=pass_sell_price,
        pass_buy_price=pass_buy_price,
        sell_price_pull=sell_price_pull,
        buy_price_pull=buy_price_pull,
        sell_price_push=sell_price_push,
        buy_price_push=buy_price_push)
    all_orders['type'] = 'order'
    all_orders['ref_buy_price'] = ref_buy_price
    all_orders['ref_sell_price'] = ref_sell_price
    all_orders['trade_edge'] = trade_edge
    all_orders['ref_edge'] = [e1, e2]
    all_orders['fill_buy_price'] = fill_buy_price
    all_orders['fill_sell_price'] = fill_sell_price
    return all_orders
